package adkr

import (
	"ebbflow/pkg/crypto"
	"encoding/json"
	"fmt"
	"log"
	"sort"
	"sync"
    "time"
)

// KeyRefreshConfig represents the configuration for key refresh
type KeyRefreshConfig struct {
	SID        string        `json:"sid"`
	PID        int           `json:"pid"`
	N          int           `json:"n"`
	F          int           `json:"f"`
	L          int           `json:"l"`
	ConfigChain [][]int      `json:"config_chain"`
	G          *crypto.G1    `json:"g"`
	Type       string        `json:"type"` // 's' for standard, 'b' for BN254
	SPK2s      []crypto.G1   `json:"spk2s"`
	SSK2       crypto.ZR     `json:"ssk2"`
	EPKs       []crypto.G1   `json:"epks"`
	ESK        crypto.ZR     `json:"esk"`
	THSKOld    crypto.ZR     `json:"thsk_old"`
	THPKOld    crypto.G1     `json:"thpk_old"`
	THPKsOld   []crypto.G1   `json:"thpks_old"`
}

// KeyRefreshOutput represents the output of key refresh protocol
type KeyRefreshOutput struct {
	NewTHPK    crypto.G1     `json:"new_thpk"`
	NewTHPKs   []crypto.G1   `json:"new_thpks"`
    NewTHSK    crypto.ZR     `json:"new_thsk"`
	ConfigChain [][]int      `json:"config_chain"`
	Sigma      crypto.G1     `json:"sigma"`
	Sigma2     crypto.G1     `json:"sigma2"`
	Digest2    []byte        `json:"digest2"`
    // Optional aggregated signature (BLS G2) over new config digest, and signers used
    AggSig     *crypto.G2    `json:"agg_sig,omitempty"`
    AggSigners []int         `json:"agg_signers,omitempty"`
}

// KeyRefreshMessage represents messages in key refresh protocol
type KeyRefreshMessage struct {
	Tag     string `json:"tag"`
	Round   int    `json:"round"`
	Payload []byte `json:"payload"`
}

// KeyRefreshTag represents different message types in key refresh
type KeyRefreshTag string

const (
	ADKR_CONFIG KeyRefreshTag = "ADKR_CONFIG"
)

// Signature share payload for key refresh config digest
type KRSigSharePayload struct {
    Round   int        `json:"round"`
    Signer  int        `json:"signer"`
    SigShare crypto.G2 `json:"sig_share"`
}

// KeyRefreshProtocol implements the key refresh protocol
type KeyRefreshProtocol struct {
	config     KeyRefreshConfig
	msgInCh    <-chan Message
	msgOutCh   chan<- Message
	outputCh   chan<- KeyRefreshOutput
	logger     *log.Logger
	mu         sync.RWMutex
	acssOutputs map[int]ACSSOutput
	completedDealers map[int]struct{}
    // new THSK shares for all nodes in new config (computed locally via interpolation)
    newTHSKs []crypto.ZR
}

// NewKeyRefreshProtocol creates a new key refresh protocol instance
func NewKeyRefreshProtocol(
	config KeyRefreshConfig,
	msgInCh <-chan Message,
	msgOutCh chan<- Message,
	outputCh chan<- KeyRefreshOutput,
	logger *log.Logger,
) *KeyRefreshProtocol {
	return &KeyRefreshProtocol{
		config:           config,
		msgInCh:          msgInCh,
		msgOutCh:         msgOutCh,
		outputCh:         outputCh,
		logger:           logger,
		acssOutputs:      make(map[int]ACSSOutput),
		completedDealers: make(map[int]struct{}),
	}
}

// Run executes the key refresh protocol
func (krp *KeyRefreshProtocol) Run() error {
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Starting key refresh protocol", krp.config.SID, krp.config.PID)
	
	// Run multiple rounds of key refresh
	for round := 0; round < len(krp.config.ConfigChain)-1; round++ {
		krp.logger.Printf("KeyRefresh-SID:%s PID:%d Starting round %d", krp.config.SID, krp.config.PID, round)
		
		output, err := krp.runRound(round)
		if err != nil {
			return fmt.Errorf("round %d failed: %w", round, err)
		}
		
		// Update configuration for next round
		krp.updateConfigForNextRound(round, output)
		
		krp.logger.Printf("KeyRefresh-SID:%s PID:%d Completed round %d", krp.config.SID, krp.config.PID, round)
	}
	
	// Send final output
	finalOutput := krp.generateFinalOutput()
	krp.outputCh <- finalOutput
	
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Key refresh protocol completed", krp.config.SID, krp.config.PID)
	return nil
}

// runRound executes a single round of key refresh
func (krp *KeyRefreshProtocol) runRound(round int) (*KeyRefreshOutput, error) {
	// Get current and next configurations
	C_o := krp.config.ConfigChain[round]
	C_n := krp.config.ConfigChain[round+1]
	N_o := len(C_o)
	N_n := len(C_n)
	
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Round %d: Old config %v, New config %v", 
		krp.config.SID, krp.config.PID, round, C_o, C_n)
	
	// Run ACSS instances for all parties in old configuration
	acssOutputs, err := krp.runACSSInstances(round, C_o, C_n, N_o, N_n)
	if err != nil {
		return nil, fmt.Errorf("ACSS instances failed: %w", err)
	}
	
	// Run MVBA to agree on which ACSS outputs to use
	selectedOutputs, err := krp.runMVBA(round, acssOutputs, C_o, N_o)
	if err != nil {
		return nil, fmt.Errorf("MVBA failed: %w", err)
	}
	
    // Compute new key shares from selected ACSS outputs (interpolation-based)
    newKeyShares, err := krp.computeNewKeyShares(selectedOutputs, acssOutputs, C_n, N_n, N_o)
	if err != nil {
		return nil, fmt.Errorf("key share computation failed: %w", err)
	}
	
	// Generate and broadcast configuration
	configOutput, err := krp.generateAndBroadcastConfig(round, newKeyShares, C_n, N_o)
	if err != nil {
		return nil, fmt.Errorf("config generation failed: %w", err)
	}
	
	return configOutput, nil
}

// runACSSInstances runs ACSS instances for all parties
func (krp *KeyRefreshProtocol) runACSSInstances(round int, C_o, C_n []int, N_o, N_n int) (map[int]ACSSOutput, error) {
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Running %d ACSS instances", krp.config.SID, krp.config.PID, N_o)
	
	// Create channels for ACSS communication
	acssInChs := make([]chan Message, N_o)
	acssOutputChs := make([]chan ACSSOutput, N_o)
	acssInputChs := make([]chan [2]crypto.ZR, N_o)
	
	for i := 0; i < N_o; i++ {
		acssInChs[i] = make(chan Message, 100)
		acssOutputChs[i] = make(chan ACSSOutput, 1)
		acssInputChs[i] = make(chan [2]crypto.ZR, 1)
	}
	
	// Create internal message router for ACSS instances
	acssInternalOutCh := make(chan Message, 1000)
	
	// Start message router for ACSS
	go func() {
		for {
			select {
			case msg := <-krp.msgInCh:
				if msg.Tag == ADKR_ACSS {
					krp.logger.Printf("KeyRefresh-SID:%s PID:%d Routing external ACSS message: Round=%d, To=%d, Sender=%d", 
						krp.config.SID, krp.config.PID, msg.Round, msg.To, msg.Sender)
					
					// Route to all ACSS instances since ECHO/READY messages can come from any node
					// and need to be processed by all ACSS instances
					routedCount := 0
					for i := 0; i < N_o; i++ {
						select {
						case acssInChs[i] <- msg:
							routedCount++
						default:
							// Channel full, skip message
							krp.logger.Printf("KeyRefresh-SID:%s PID:%d ACSS channel %d full, skipping message", 
								krp.config.SID, krp.config.PID, i)
						}
					}
					krp.logger.Printf("KeyRefresh-SID:%s PID:%d Routed external ACSS message to %d/%d instances", 
						krp.config.SID, krp.config.PID, routedCount, N_o)
				}
			case msg := <-acssInternalOutCh:
				if msg.Tag == ADKR_ACSS {
					// Parse ACSS message to get more details
					var acssMsg ACSSMessage
					if err := json.Unmarshal(msg.Payload, &acssMsg); err == nil {
						krp.logger.Printf("KeyRefresh-SID:%s PID:%d Routing internal ACSS message: Tag=%s, Round=%d, To=%d, Sender=%d", 
							krp.config.SID, krp.config.PID, acssMsg.Tag, msg.Round, msg.To, msg.Sender)
					} else {
						krp.logger.Printf("KeyRefresh-SID:%s PID:%d Routing internal ACSS message: Round=%d, To=%d, Sender=%d", 
							krp.config.SID, krp.config.PID, msg.Round, msg.To, msg.Sender)
					}
					
					// Route internal ACSS messages to all instances
					routedCount := 0
					for i := 0; i < N_o; i++ {
						select {
						case acssInChs[i] <- msg:
							routedCount++
						default:
							// Channel full, skip message
							krp.logger.Printf("KeyRefresh-SID:%s PID:%d ACSS channel %d full, skipping internal message", 
								krp.config.SID, krp.config.PID, i)
						}
					}
					krp.logger.Printf("KeyRefresh-SID:%s PID:%d Routed internal ACSS message to %d/%d instances", 
						krp.config.SID, krp.config.PID, routedCount, N_o)
					
					// Also forward to external output
					krp.msgOutCh <- msg
				}
			}
		}
	}()
	
	// Start ACSS instances
	var wg sync.WaitGroup
	for i := 0; i < N_o; i++ {
		wg.Add(1)
		go func(dealerID int) {
			defer wg.Done()
			
			var inputCh <-chan [2]crypto.ZR
			if C_o[dealerID] == krp.config.PID {
				inputCh = acssInputChs[dealerID]
			}
			
			CompleteSecretSharing(
				krp.config.SID+"ACSS"+fmt.Sprintf("%d-%d", round, dealerID),
				krp.config.PID, N_o, krp.config.F, C_o, N_n, krp.config.F, C_n,
				krp.config.G, krp.config.G, dealerID,
				krp.config.EPKs, krp.config.ESK,
				inputCh, acssOutputChs[dealerID], acssInChs[dealerID], acssInternalOutCh,
			)
		}(i)
	}
	
	// Provide input for our own ACSS instance
	if krp.isInOldConfig(round) {
		ourIndex := krp.getIndexInOldConfig(round)
		secretA := crypto.ZRFromInt(int64(123 + krp.config.PID + round))
		secretB := crypto.ZRFromInt(int64(456 + krp.config.PID + round))
		acssInputChs[ourIndex] <- [2]crypto.ZR{secretA, secretB}
	}
	
	// Collect ACSS outputs with timeout
	acssOutputs := make(map[int]ACSSOutput)
	acssTimeout := func() time.Duration {
		if IsFastBench() { return 30 * time.Second }
		return 60 * time.Second
	}()
	
	acssDeadline := time.Now().Add(acssTimeout)
	timeoutReached := false
	
	for i := 0; i < N_o && !timeoutReached; i++ {
		select {
		case output := <-acssOutputChs[i]:
			acssOutputs[i] = output
			krp.logger.Printf("KeyRefresh-SID:%s PID:%d Received ACSS output from dealer %d", 
				krp.config.SID, krp.config.PID, i)
		case <-time.After(time.Until(acssDeadline)):
			krp.logger.Printf("KeyRefresh-SID:%s PID:%d ACSS collection timeout, got %d/%d outputs", 
				krp.config.SID, krp.config.PID, len(acssOutputs), N_o)
			timeoutReached = true
		}
	}
	
	// If we have at least f+1 outputs, we can proceed
	if len(acssOutputs) < krp.config.F+1 {
		return nil, fmt.Errorf("insufficient ACSS outputs: got %d, need at least %d", len(acssOutputs), krp.config.F+1)
	}
	
	wg.Wait()
	return acssOutputs, nil
}

// runMVBA runs MVBA to agree on which ACSS outputs to use
func (krp *KeyRefreshProtocol) runMVBA(round int, acssOutputs map[int]ACSSOutput, C_o []int, N_o int) ([]int, error) {
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Running MVBA for round %d", krp.config.SID, krp.config.PID, round)
    // Guard: if there are no parties in the old configuration, skip
    if N_o == 0 {
        return []int{}, nil
    }

    // Branch on SPBC fast-path vs MVBA
    if IsKeyRefreshSPBCFast() {
        // SPBC fast-path: leader proposes a set, all adopt after SPBC deliver
        leaderPID := C_o[0]
        proposalSet := krp.selectProposalSet(acssOutputs)
        proposalBytes, _ := json.Marshal(proposalSet)
        spbcVal := SPBCVal{RootHash: []byte{}, Branch: [][]byte{}, Stripe: proposalBytes}
        spbcValBytes, _ := json.Marshal(spbcVal)

        deliveredCh := make(chan []byte, 1)
        go StrongProvableBroadcast(
            krp.config.SID+"SPBC"+fmt.Sprintf("%d", round),
            krp.config.PID,
            N_o, krp.config.F,
            leaderPID,
            round,
            func() []byte { if krp.config.PID == leaderPID { return spbcValBytes }; return nil }(),
            krp.msgInCh,
            krp.msgOutCh,
            func(val SPBCVal) { deliveredCh <- val.Stripe },
        )

        select {
        case decided := <-deliveredCh:
            var selectedDealers []int
            if err := json.Unmarshal(decided, &selectedDealers); err != nil {
                return nil, fmt.Errorf("failed to unmarshal SPBC decision: %w", err)
            }
            krp.logger.Printf("KeyRefresh-SID:%s PID:%d SPBC decided on dealers: %v", krp.config.SID, krp.config.PID, selectedDealers)
            // Experimental fast-bench: trigger a simplified coin run to emit metrics
            if IsFastBench() {
                coinSID := krp.config.SID + "-coin"
                // use registered or fallback keys
                pks, thpk := GetThresholdBLSPKs()
                var mySK crypto.ZR
                if GetMyBLSSK() != nil { mySK = *GetMyBLSSK() } else {
                    genPK, _, sks := crypto.GenerateThresholdKeys(N_o, krp.config.F+1)
                    thpk = genPK; mySK = sks[krp.config.PID]
                }
                tcp := NewThresholdCoinProtocol(coinSID, krp.config.PID, N_o, krp.config.F, 0, C_o, krp.config.G, 0, pks, mySK, thpk, krp.msgInCh, krp.msgOutCh, log.Default())
                go func(){ _, _ = tcp.Run(0) }()
            }
            return selectedDealers, nil
        case <-time.After(func() time.Duration {
            if IsFastBench() { return 5 * time.Second }
            return 30 * time.Second
        }()):
            return nil, fmt.Errorf("SPBC timed out")
        }
    }

    // MVBA path (default): start N RBCs and run MVBA+ABA+coin on proposal sets
    // Prepare MVBA and coin channels
    mvbaInCh := make(chan Message, 100)
    coinInCh := make(chan Message, 100)
    // Dispatch messages to MVBA and coin
    go func() {
        for msg := range krp.msgInCh {
            if msg.Tag == ADKR_MVBA { 
                mvbaInCh <- msg 
            } else if msg.Tag == ADKG_COIN {
                coinInCh <- msg
            }
        }
    }()

    mvbaInputCh := make(chan []byte, 1)
    mvbaOutputCh := make(chan []byte, 1)
    localSetCh := make(chan int, N_o)

    // Local set: ACSS completed dealers
    for dealerID := range acssOutputs { localSetCh <- dealerID }
    close(localSetCh)

    predicate := DefaultMVBAPredicate(N_o, krp.config.F)

    // Threshold BLS for coin: use registered materials or fallback generate
    pks, thpk := GetThresholdBLSPKs()
    var mySK crypto.ZR
    if GetMyBLSSK() != nil { mySK = *GetMyBLSSK() } else {
        genPK, _, sks := crypto.GenerateThresholdKeys(N_o, krp.config.F+1)
        thpk = genPK; mySK = sks[krp.config.PID]
    }

    go TThresholdMVBA(
        krp.config.SID+"MVBA"+fmt.Sprintf("%d", round),
        krp.config.PID, N_o, krp.config.F, C_o, krp.config.G,
        mvbaInputCh, mvbaOutputCh, mvbaInCh, krp.msgOutCh,
        nil, localSetCh, predicate, pks, mySK, thpk,
    )

    // Propose our local set
    proposalSet := krp.selectProposalSet(acssOutputs)
    proposalBytes, _ := json.Marshal(proposalSet)
    mvbaInputCh <- proposalBytes

    select {
    case decided := <-mvbaOutputCh:
        var selectedDealers []int
        if err := json.Unmarshal(decided, &selectedDealers); err != nil { return nil, fmt.Errorf("MVBA decide unmarshal: %w", err) }
        krp.logger.Printf("KeyRefresh-SID:%s PID:%d MVBA decided dealers: %v", krp.config.SID, krp.config.PID, selectedDealers)
        return selectedDealers, nil
    case <-time.After(func() time.Duration {
        if IsFastBench() { return 10 * time.Second }
        return 60 * time.Second
    }()):
        return nil, fmt.Errorf("MVBA timed out")
    }
}

// computeNewKeyShares computes new key shares from selected ACSS outputs
func (krp *KeyRefreshProtocol) computeNewKeyShares(selectedDealers []int, acssOutputs map[int]ACSSOutput, C_n []int, N_n int, N_o int) ([]crypto.G1, error) {
    krp.logger.Printf("KeyRefresh-SID:%s PID:%d Computing new key shares via interpolation", krp.config.SID, krp.config.PID)

    // Build (x,y) pairs from selected dealers: x = dealer index + 1, y = ShareA (as ZR)
    coords := make([][2]interface{}, 0, len(selectedDealers))
    for _, dealerID := range selectedDealers {
        out, ok := acssOutputs[dealerID]
        if !ok { continue }
        x := int64(dealerID + 1)
        y := out.ShareA
        coords = append(coords, [2]interface{}{x, y})
    }

    newKeyShares := make([]crypto.G1, N_n)
    krp.newTHSKs = make([]crypto.ZR, N_n)
    for i := 0; i < N_n; i++ {
        // Interpolate polynomial at x = i+1 over selected dealer shares
        xi := int64(i + 1)
        // crypto.InterpolateAtX expects slice of pairs and returns ZR
        zi, err := crypto.InterpolateAtX(coords, xi)
        if err != nil {
            // fallback: zero share
            zi.SetZero()
        }
        krp.newTHSKs[i] = zi
        newKeyShares[i] = crypto.MulG1(krp.config.G, &zi)
        if C_n[i] == krp.config.PID {
            krp.config.THSKOld = zi
        }
    }
    return newKeyShares, nil
}

// generateAndBroadcastConfig generates and broadcasts the new configuration
func (krp *KeyRefreshProtocol) generateAndBroadcastConfig(round int, newKeyShares []crypto.G1, C_n []int, N_o int) (*KeyRefreshOutput, error) {
	krp.logger.Printf("KeyRefresh-SID:%s PID:%d Generating and broadcasting config", krp.config.SID, krp.config.PID)
	
	// Compute new threshold public key
	newTHPK := krp.computeNewThresholdPublicKey(newKeyShares)

    // Edge case: no parties in old config, skip signature aggregation
    if N_o == 0 {
        return &KeyRefreshOutput{
            NewTHPK:     newTHPK,
            NewTHPKs:    newKeyShares,
            ConfigChain: [][]int{C_n},
        }, nil
    }
	
    // Generate digest for config
    digest := crypto.HashToG1([]byte(fmt.Sprintf("%v", newTHPK)))
    digest2 := crypto.HashToG1([]byte(fmt.Sprintf("%v%v", newTHPK, C_n)))
	
    // Threshold signature over digest2 using NEW THSK shares derived above
    signerSet := make([]int, 0, krp.config.F+1)
    sigShares := make(map[int]*crypto.G2)
    // pick first f+1 indices in NEW config
    for i := 0; i < len(C_n) && len(signerSet) < krp.config.F+1; i++ {
        signerSet = append(signerSet, i)
        skShare := krp.newTHSKs[i]
        sigShares[i] = crypto.Sign(&skShare, digest2.G1Affine.Bytes())
    }
    agg := crypto.CombineShares(sigShares, signerSet)
    // Verify aggregate signature against NEW threshold PK
    _ = crypto.VerifySignature(&newTHPK, agg, digest2.G1Affine.Bytes())
	
	// Convert digest2 to bytes
	digest2Bytes := digest2.G1Affine.Bytes()
	
	// Create configuration message
	configMsg := KeyRefreshMessage{
		Tag:   string(ADKR_CONFIG),
		Round: round,
		Payload: []byte(fmt.Sprintf("config_%d", round)),
	}
	
	// Broadcast configuration
	for i := 0; i < N_o; i++ {
		payloadBytes, _ := json.Marshal(configMsg)
		krp.msgOutCh <- Message{
			Tag:     BroadcastTag(ADKR_CONFIG),
			Round:   round,
			Sender:  krp.config.PID,
			To:      i,
			Payload: payloadBytes,
		}
	}
	
    return &KeyRefreshOutput{
		NewTHPK:     newTHPK,
		NewTHPKs:    newKeyShares,
		ConfigChain: [][]int{C_n},
        Sigma:       crypto.MulG1(&digest, &krp.config.THSKOld),
        Sigma2:      crypto.MulG1(&digest2, &krp.config.THSKOld),
		Digest2:     digest2Bytes[:],
        AggSig:      agg,
        AggSigners:  signerSet,
	}, nil
}

// Helper functions

func (krp *KeyRefreshProtocol) isInOldConfig(round int) bool {
	if round >= len(krp.config.ConfigChain) {
		return false
	}
	for _, pid := range krp.config.ConfigChain[round] {
		if pid == krp.config.PID {
			return true
		}
	}
	return false
}

func (krp *KeyRefreshProtocol) getIndexInOldConfig(round int) int {
	if round >= len(krp.config.ConfigChain) {
		return -1
	}
	for i, pid := range krp.config.ConfigChain[round] {
		if pid == krp.config.PID {
			return i
		}
	}
	return -1
}

func (krp *KeyRefreshProtocol) selectProposalSet(acssOutputs map[int]ACSSOutput) []int {
	// Select first f+1 dealers that have completed ACSS
	var selected []int
	for dealerID := range acssOutputs {
		if len(selected) >= krp.config.F+1 {
			break
		}
		selected = append(selected, dealerID)
	}
	sort.Ints(selected)
	return selected
}

func (krp *KeyRefreshProtocol) computeNewThresholdPublicKey(newKeyShares []crypto.G1) crypto.G1 {
	// Simplified implementation - in practice this would involve proper interpolation
	var result crypto.G1
	result.SetInfinity()
	
	for _, share := range newKeyShares {
		result = crypto.AddG1(&result, &share)
	}
	
	return result
}

func (krp *KeyRefreshProtocol) updateConfigForNextRound(round int, output *KeyRefreshOutput) {
	// Update configuration for next round
	krp.mu.Lock()
	defer krp.mu.Unlock()
	
	// Update threshold keys
	krp.config.THPKOld = output.NewTHPK
	krp.config.THPKsOld = output.NewTHPKs
	// Note: THSKOld would need to be updated based on the new shares
}

func (krp *KeyRefreshProtocol) generateFinalOutput() KeyRefreshOutput {
	krp.mu.RLock()
	defer krp.mu.RUnlock()
	
	return KeyRefreshOutput{
		NewTHPK:     krp.config.THPKOld,
		NewTHPKs:    krp.config.THPKsOld,
		ConfigChain: krp.config.ConfigChain,
	}
}
